# -*- coding: utf-8 -*-
# file: lstm.py
# author: songyouwei <youwei0314@gmail.com>
# Copyright (C) 2018. All Rights Reserved.

from layers.dynamic_rnn import DynamicLSTM
import torch
import torch.nn as nn


class LSTM(nn.Module):
    def __init__(self, embedding_matrix, opt):
        super(LSTM, self).__init__()
        self.embed = nn.Embedding.from_pretrained(torch.tensor(embedding_matrix, dtype=torch.float))   # 传入词嵌入矩阵，输入的序列中的token原来是词典中的
        #inputsize,hiddensize,LSTM层数（如果是两层，就是第二层LSTM接受第一层的输出作为输入），是否使用偏置
        self.lstm = DynamicLSTM(opt.embed_dim, opt.hidden_dim, num_layers=1, batch_first=True)  #(batch, seq, feature)
        self.dense = nn.Linear(opt.hidden_dim, opt.polarities_dim)  #直接接上一个全连接

    def forward(self, inputs):  # lstm模型的输入的就是文本的index，inputs.shape=[bs,seq_len],LSTM所需要的输入为text_indices
        text_raw_indices = inputs[0]  # inputs为一个List，拿到List中的tensor,tensor.shape=(16,85)
        x = self.embed(text_raw_indices)  # 转换为词嵌入，得到的x就是一个词嵌入矩阵了,压缩输入也是压缩embed之后的输入
        x_len = torch.sum(text_raw_indices != 0, dim=-1)  # text_rw_indeices = tensor(16,85), 在-1维度计算非0元素个数得到batch中每个句子的长度，x_len = [16] 得到一个batch所包含的句子长度？
        _, (h_n, _) = self.lstm(x, x_len)  #输入的x为(N,L,Hin）输出包含output,h_n,c_n
        out = self.dense(h_n[0])   #h_n=(D*num_layers,N,Hout),h_n[0]=(N,Hout),h_n[0]代表正向的，Hout是隐藏向量的输出，axis=0时的每个元素就是一个句子经过处理之后得到的最终向量
        return out  #返回的为极性的概率


'''
可以挤压掉一个维度
     h_n=torch.squeeze(h_n)
    out = self.dense(h_n)   #h_n=(D*num_layers,N,Hout),h_n[0]=(N,Hout),h_n[0]代表正向的，Hout是隐藏向量的输出，axis=0时的每个元素就是一个句子经过处理之后得到的最终向量
'''
